'''
Pydantic 深度解析：

Pydantic 的核心价值 = 数据模型定义 + 自动类型校验 + 数据转换 + 序列化/反序列化

想象一下：你点外卖时，期望收到的是「汉堡+薯条+可乐」，但实际可能收到：
- 格式不对：送来了一堆散装食材 ❌
- 类型不对：汉堡变成了包子 ❌
- 缺少内容：忘了给可乐 ❌

Pydantic 就是确保你收到的数据完全符合预期格式的「质检员」！
'''

import os
from dotenv import load_dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

load_dotenv()

# ========== PYDANTIC 数据模型定义 ==========
class Joke(BaseModel):
    """
    Pydantic数据模型：定义我们期望从大模型获得的数据结构

    为什么要用Pydantic？
    1. 明确约束：告诉大模型必须返回什么字段，什么类型
    2. 自动验证：如果返回的数据不符合这个结构，Pydantic会报错
    3. 类型安全：确保setup是字符串，punchline也是字符串
    4. 文档化：Field的description会成为给大模型的指令的一部分
    """
    setup: str = Field(description="设置笑话的问题部分", example="为什么程序员不喜欢大自然？")
    punchline: str = Field(description="解决笑话的答案/笑点部分", example="因为里面有太多bug！")

# 如果没有Pydantic，大模型可能返回：
# - 不完整的JSON：只有setup没有punchline ❌
# - 错误类型：setup返回成了数字 ❌
# - 额外字段：多了一个没用的字段 ❌
# - 格式混乱：根本不是JSON ❌

# ========== 大模型配置 ==========
llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),
    model="deepseek-ai/DeepSeek-R1-0528",
    streaming=True,
    temperature=0.7
)

joke_query = "告诉我一个笑话。"

# ========== PYDANTIC + JSON解析器 ==========
# 关键步骤：将Pydantic模型传递给解析器
parser = JsonOutputParser(pydantic_object=Joke)

# 查看自动生成的格式指令
print("=== 自动生成的格式指令 ===")
format_instructions = parser.get_format_instructions()
print(format_instructions)
print("=" * 50)

'''
parser.get_format_instructions() 会自动生成类似这样的指令：
"The output should be formatted as a JSON instance that conforms to the JSON schema below."

这样大模型就知道：
- 必须返回JSON格式
- 必须包含setup和punchline两个字段  
- 两个字段都必须是字符串
- 字段的含义是什么
'''

prompt = PromptTemplate(
    template="回答用户的查询。\n{format_instructions}\n{query}\n",
    input_variables=["query"],
    partial_variables={"format_instructions": format_instructions}
)

# ========== 创建处理链 ==========
chain = prompt | llm | parser

# ========== 测试验证 ==========
print("=== 测试Pydantic验证效果 ===")
response = chain.invoke({"query": joke_query})
print("完整响应:", response)
print("数据类型:", type(response))  # 应该是dict
print("setup字段:", response['setup'])
print("punchline字段:", response['punchline'])

# ========== 手动演示Pydantic验证 ==========
print("\n=== 手动演示Pydantic验证 ===")

# 正确的数据
good_data = {"setup": "test", "punchline": "test"}
joke_instance = Joke(**good_data)
print("✅ 正确数据验证通过:", joke_instance)

# 错误的数据演示
try:
    bad_data = {"setup": "test"}  # 缺少punchline字段
    joke_instance_bad = Joke(**bad_data)
except Exception as e:
    print("❌ 错误数据被捕获:", e)

# ========== 流式输出 ==========
print("\n\n\n=== 流式输出演示 ===")
print("注意：即使流式输出，每个片段也是符合Joke结构的JSON")
for chunk in chain.stream({"query": joke_query}):
    print("流式片段:", chunk)